from strategy.core.Stg import Stg
import talib as ta
import math
import pandas as pd


class BollStg(Stg):
    preBuySell = 0
    param = {
        'timeperiod': 5,
        'nbdev': 2,  # 多少倍标准差
        'type': 'press_support'  # press_support: 压力支撑, trend: 趋势
    }

    def initial(self):
        self.upperband, self.middleband, self.lowerband = \
            ta.BBANDS(self.data['close'], nbdevup=float(self.param['nbdev']), nbdevdn=float(self.param['nbdev']),
                      timeperiod=self.param['timeperiod'], matype=0)
        return self.upperband, self.middleband, self.lowerband

    def signal(self, index):
        # 算法
        # k或d<20: 超卖, 应买入
        # k或d>80: 超买, 应卖出
        # 自动跳过无法计算的节点
        if index < self.param['timeperiod']:
            return
        pre = index - 1
        res = 0
        if self.param['type'] == 'pressSupport':
            # 使用支持和阻力理论, 股价到上轨(upperband)有阻力卖出,股价到下轨(lowerband)有阻力买入, 但该理论与趋势理论相违背
            if self.data['close'][index] >= self.upperband[index]:
                return self.sell()
            elif self.data['close'][index] <= self.lowerband[index]:
                return self.buy()
        elif self.param['type'] == 'trend':
            # 使用趋势理论, 股价上穿上轨(upperband)买入, 股价下穿下轨(lowerband)卖出, 但该理论与支持和阻力理论相违背
            if self.data['close'][pre] < self.upperband[pre] and self.data['close'][index] >= self.upperband[index]:
                return self.buy()
            elif self.data['close'][pre] > self.lowerband[pre] and self.data['close'][index] <= self.lowerband[index]:
                return self.sell()
        elif self.param['type'] == 'openTrade':
            # 使用开口加趋势理论, 上下轨开口越大+股价超过上轨买跌破下轨买, 趋势越强, 但该理论与支持和阻力理论相违背
            if self.upperband[pre] - self.lowerband[pre] < self.upperband[index] - self.lowerband[index]:
                if self.data['close'][index] >= self.upperband[index]:
                    return self.buy()
                elif self.data['close'][index] <= self.lowerband[index]:
                    return self.sell()

    def buy(self):  # 在上次买后不再买, 防止重复出现买状态
        if self.preBuySell != 1:
            self.preBuySell = 1
            return 1

    def sell(self):
        if self.preBuySell != -1:
            self.preBuySell = -1
            return -1
